using System;
using System.Collections.Generic;
using System.IO;

namespace DeflateLib
{
    /// <summary>
    /// DEFLATE decoder algorithm implementation.
    /// </summary>
    sealed class DecoderWithInitialization
    {
		InfiniteBuffer buffer = new InfiniteBuffer(DeflateUtils.MaxDistance + 1);

		Stream input;
        BitReader reader;
		
        /// <summary>
        /// Gets base stream that contains compressed data. 
        /// Was used in constructor.
        /// </summary>
        public Stream Input
        {
			get { return input; }
        }

        /// <summary>
        /// Creates instance of the decoder.
        /// </summary>
        /// <param name="input">Stream that contains compressed data.</param>
        public DecoderWithInitialization(Stream input)
		{
            if (input == null) throw new ArgumentNullException("input");

			this.input = input;
            this.reader = new BitReader(input);
        }

        /// <summary>
        /// Sets dictionary for decoder.
        /// </summary>
        /// <param name="data">Dictionary data</param>
        /// <param name="index">Offset in data</param>
        /// <param name="length">Data length</param>
        public void SetDictionary(byte[] data, int index, int length)
        {
            if (data == null) throw new ArgumentNullException("data");
            if (index < 0 || length < 0 || index + length > data.Length)
                throw new ArgumentException("index or length is outside data");

            buffer.Write(data, index, length);
        }

        internal void SetInitializationData(InitializationData data)
        {
            if (data == null) throw new ArgumentNullException("data");

            this.blockType = data.BlockType;
            this.isFinalBlock = data.IsFinal;
            this.isEobFound = false;

            switch(data.BlockType)            
            {
                case DeflateUtils.NoCompresstionBlockType:
                    this.block0DataLeft = data.Block0Limit;
                    this.huffmanTrees = null;
                    break;
                case DeflateUtils.StandardCodeBlockType:
                    this.block0DataLeft = 0;
                    this.huffmanTrees = DeflateUtils.StaticCodesAndDistances;
                    break;
                case DeflateUtils.DynamicCodeBlockType:
                    this.block0DataLeft = 0;
                    this.huffmanTrees = data.Block2HuffmanTrees;
                    break;
            }

            reader.SetBitPosition(data.InputOffset, data.BitPosition);            
        }

        internal CodesAndDistances ReadDynamicTrees(long inputOffset, long treePosition)
        {
            reader.SetBitPosition(inputOffset, treePosition);
            TreeNode[] codes, distances;
            ReadDynamic(reader, out codes, out distances);
            return DeflateUtils.BuildHuffmanTrees(codes, distances);
        }

        private void OnUnefficientTree()
        {
        }

        int blockType = -1;
        bool isFinalBlock = false;
        int block0DataLeft = 0;
        long blockBitPosition = 0;
        long lastBitPosition = 0;

        CodesAndDistances huffmanTrees;

        bool endOfBlocks = false;

        /// <summary>
        /// Reads decoded data. Returns all read data or to first sync event.
        /// </summary>
        /// <param name="data">Result data buffer</param>
        /// <param name="index"></param>
        /// <param name="length"></param>
        /// <returns></returns>
        public int Read(byte[] data, int index, int length)
        {
            if (data == null) throw new ArgumentNullException("data");
            if (index < 0 || length < 0 || index + length > data.Length)
                throw new ArgumentException("index or length is outside data");

            if (length == 0 || endOfBlocks) return 0;

            int position = 0;
            while (position < length)
            {
                EnsureHeaderRead();

                if (endOfBlocks) break;

                int read = ReadBlockData(data, index + position, length - position);
                if (read > 0)
                    position += read;
                else
                    blockType = -1;
            }
            return position;
        }

        private void EnsureHeaderRead()
        {
            while (IsBlockHeaderExpected() && !endOfBlocks)
            {
                endOfBlocks = !ReadBlockHeader();
            }
        }

        private bool IsBlockHeaderExpected()
        {
            return blockType < 0;
        }

        private bool ReadBlockHeader()
        {
            if (this.isFinalBlock)
                return false;

            blockBitPosition = reader.CurrentPosition;

            this.isFinalBlock = reader.ReadBit();
            this.blockType = reader.ReadLSB(2);

            TreeNode[] codes;
            TreeNode[] distances;

            switch (blockType)
            {
                case DeflateUtils.NoCompresstionBlockType:
                    reader.Align();
                    int len = reader.ReadLSB(16);
                    int nlen = reader.ReadLSB(16);
                    if ((len & ~nlen) != len) throw new DeflateLibException("Invalid length of no compression block");

                    this.block0DataLeft = len;
                    this.huffmanTrees = null;
                    this.isEobFound = true;
                    break;
                case DeflateUtils.StandardCodeBlockType:
                    codes = DeflateUtils.StaticCodes;
                    distances = DeflateUtils.StaticDistanceCodes;

                    this.block0DataLeft = 0;
                    this.huffmanTrees = DeflateUtils.StaticCodesAndDistances;
                    this.isEobFound = false;
                    break;
                case DeflateUtils.DynamicCodeBlockType:
                    ReadDynamic(reader, out codes, out distances);

                    this.block0DataLeft = 0;
                    this.huffmanTrees = DeflateUtils.BuildHuffmanTrees(codes, distances);
                    this.isEobFound = false;
                    break;
                default:
                    throw new DeflateLibException("Unknown block type 3");
            }

            this.lastBitPosition = reader.CurrentPosition;
            return true;
        }

        int delayedOffset = -1;
        int delayedCount = 0;
        bool isEobFound = false;

        private int ReadBlockData(byte[] data, int offset, int count)
        {
            int startOffset = offset;
            if (blockType == DeflateUtils.NoCompresstionBlockType)
            {
                if (block0DataLeft > 0)
                {
                    int countToGet = Math.Min(count, block0DataLeft);
                    reader.GetBytes(data, offset, countToGet);
                    buffer.Write(data, offset, countToGet);
                    block0DataLeft -= countToGet;
                    count -= countToGet;
                }
            }
            else if(!isEobFound)
            {                
                if (delayedCount > 0)
                {
                    ReadBufferedData(data, ref offset, ref count);
                }

                if (count > 0)
                {
                    do
                    {
                        int code = ReadHuffmanSequence(reader, huffmanTrees.Codes);
                        isEobFound = code == DeflateUtils.EndOfBlockCode;
                        if (isEobFound) break;
                        if (code < DeflateUtils.EndOfBlockCode)
                        {
                            data[offset++] = (byte)code;
                            buffer.WriteByte((byte)code);
                            count--;
                        }
                        else if (code <= DeflateUtils.MaxCode)
                        {
                            int length = DecodeLength(reader, code);
                            int distance = ReadAndDecodeDistance(reader, huffmanTrees.Distances);

                            if (distance > buffer.RealPosition)
                                throw new DeflateLibException("Distance cannot refer past the beginning");

                            delayedOffset = distance;
                            delayedCount = length;

                            ReadBufferedData(data, ref offset, ref count);
                        }
                    } while (count > 0);

                }
            }

            lastBitPosition = reader.CurrentPosition;
            return offset - startOffset;
        }

        private void ReadBufferedData(byte[] data, ref int offset, ref int count)
        {
            int countToGet = Math.Min(delayedCount, count);
            byte[] bufferData = buffer.GetData(delayedOffset,
                Math.Min(countToGet, delayedOffset));

            count -= countToGet;
            delayedCount -= countToGet;

            while (countToGet > bufferData.Length)
            {
                Array.Copy(bufferData, 0, data, offset, bufferData.Length);
                offset += bufferData.Length;
                countToGet -= bufferData.Length;
                buffer.Write(bufferData, 0, bufferData.Length);
            }
            Array.Copy(bufferData, 0, data, offset, countToGet);
            offset += countToGet;
            buffer.Write(bufferData, 0, countToGet);
        }

        public bool Skip(int count)
        {
            const int SkipDataLength = 1024;
            byte[] skipData = new byte[SkipDataLength];
            int read;
            while (count > 0 && (read  = Read(skipData, 0, Math.Min(SkipDataLength, count))) > 0)
            {
                count -= read;
            }
            return count <= 0;
        }

        /// <summary>
        /// Skips data to the end.
        /// </summary>
        public void SkipToEnd()
        {
            const int SkipDataLength = 1024;
            byte[] skipData = new byte[SkipDataLength];
            int read;
            while ((read = Read(skipData, 0, SkipDataLength)) > 0) ;
        }

        static int ReadHuffmanSequence(BitReader reader, HuffmanTree tree)
        {
            while (tree != null && !tree.IsLeaf)
            {
                tree = reader.ReadBit() ? tree.One : tree.Zero;
            }
            if (tree == null) throw new DeflateLibException("Invalid huffman code");
            return tree.SymbolIndex;
        }
        
        static int DecodeLength(BitReader reader, int code)
        {
            int start, bitsToRead;
            DeflateUtils.DecodeLength(code, out start, out bitsToRead);
            if(bitsToRead > 0)
            {
	            return start + reader.ReadLSB(bitsToRead);
            }
            else
            {
		        return start;
            }
        }

        static int ReadAndDecodeDistance(BitReader reader, HuffmanTree distancesHuffman)
        {
            int distanceCode = ReadHuffmanSequence(reader, distancesHuffman);
            if (distanceCode > DeflateUtils.MaxDistanceCode)
                throw new DeflateLibException("Invalid distance code");

            int start, bitsToRead;
            DeflateUtils.DecodeDistance(distanceCode, out start, out bitsToRead);
            if(bitsToRead > 0)
            {
                int extra = reader.ReadLSB(bitsToRead);
                if (distanceCode == 284 && extra > 30) throw new DeflateLibException("Invalid distance code extra (284)");
				return start + extra;
            }
            else
            {
				return start;
            }
        }

        void ReadDynamic(BitReader reader, out TreeNode[] codes, out TreeNode[] distances)
        {
            int hlit = reader.ReadLSB(5) + 257; // 257..288
            int hdist = reader.ReadLSB(5) + 1; // 1..32
            int hclen = reader.ReadLSB(4) + 4; // 4..19

            if (hlit > 286) throw new DeflateLibException("Invalid tree literals amount");

            int[] map = DeflateUtils.LengthsMap;
            int[] clen = new int[DeflateUtils.MaxSpecialLengthCode + 1];
            for (int i = 0; i < hclen; i++)
            {
                clen[map[i]] = reader.ReadLSB(3);
            }

            TreeNode[] lengthsAlphabet = DeflateUtils.BuildCodes(clen);
            HuffmanTree lengthsHuffman = DeflateUtils.BuildHuffmanTree(lengthsAlphabet);
            int[] lengthsSequence = ReadLengthsSequence(reader, lengthsHuffman, hlit + hdist);

            codes = new TreeNode[hlit];
            for (int i = 0; i < hlit; i++)
            {
                codes[i].Len = lengthsSequence[i];
            }
            DeflateUtils.BuildCodes(codes);

            distances = new TreeNode[hdist];
            for (int i = 0; i < hdist; i++)
            {
                distances[i].Len = lengthsSequence[i + hlit];
            }
            DeflateUtils.BuildCodes(distances);

            DeflateUtils.DumpTrees("decoder.log", clen, codes, distances);
        }


        static int[] ReadLengthsSequence(BitReader reader, HuffmanTree lengthsHuffman, int length)
        {
            int[] result = new int[length];
            for (int i = 0; i < length; i++)
            {
                int code = ReadHuffmanSequence(reader, lengthsHuffman);

                if (code < 16)
                {
                    result[i] = code;
                }
                else if (code == 16)
                {
                    int repeat = reader.ReadLSB(2) + 3;
                    if(repeat + i > result.Length)
                        throw new DeflateLibException("Invalid code length extra bits (16)");
                    for (int q = 0; q < repeat; q++)
                    {
                        result[i + q] = result[i - 1];
                    }
                    i += repeat - 1;
                }
                else if (code == 17)
                {
                    int repeat = reader.ReadLSB(3) + 3;
                    if (repeat + i > result.Length)
                        throw new DeflateLibException("Invalid code length extra bits (17)");
                    i += repeat - 1;
                }
                else if (code == 18)
                {
                    int repeat = reader.ReadLSB(7) + 11;
                    if (repeat + i > result.Length)
                        throw new DeflateLibException("Invalid code length extra bits (18)");
                    i += repeat - 1;
                }
                else
                    throw new DeflateLibException("Invalid code length");
            }
            return result;
        }

        internal class InitializationData
        {
            public int BlockType;
            public bool IsFinal;
            public long InputOffset;
            public long BitPosition;
            public CodesAndDistances Block2HuffmanTrees;
            public int Block0Limit;
        }
    }
}
